# 导入必要的库和模块
import math  # 导入数学库，用于数学计算
from typing import Optional, Tuple  # 导入类型提示模块

import torch  # 导入PyTorch库，用于张量操作和深度学习
from torch import nn  # 从PyTorch导入神经网络模块

from .activations import ACTIVATIONS  # 从当前模块导入激活函数
from .compressive_memory import CompressiveMemory  # 导入压缩记忆模块
from .positional_embeddings import PositionEmbeddings  # 导入位置嵌入模块


class InfiniTransformer(nn.Module):
    """包含压缩记忆的Transformer层。"""

    """
    该代码段定义了一个名为 `InfiniTransformer` 的类，它初始化了一个包含多头注意力和多层感知机（MLP）的模块。该类用于处理输入维度 `dim_input`，并通过压缩记忆（CompressiveMemory）机制进行注意力计算。MLP部分使用指定的激活函数 `activation`，并支持可学习的初始状态 `init_state_learnable` 和可选的位置嵌入 `position_embedder`。此外，该类还支持不同的记忆更新规则 `update` 和因果注意力掩码 `causal`。
    
    参数：
    - dim_input (int): 输入维度。
    - dim_hidden (int): MLP的隐藏层维度。
    - dim_key (int): 压缩记忆中键的维度。
    - dim_value (int): 压缩记忆中值的维度。
    - num_heads (int): 压缩记忆中的注意力头数。
    - activation (str): MLP使用的激活函数名称。
    - segment_len (int): 压缩记忆中的段长度。
    - update (str, optional): 压缩记忆的记忆更新规则类型（"linear" 或 "delta"），默认为 "linear"。
    - causal (bool, optional): 是否使用因果注意力掩码，默认为 False。
    - position_embedder (Optional[PositionEmbeddings], optional): 位置嵌入模块，默认为 None。
    - init_state_learnable (bool, optional): 压缩记忆的初始状态是否可学习，默认为 False。
    - dropout (float, optional): MLP的丢弃率，默认为 0.0。
    """

    def __init__(
        self,
        dim_input: int,  # 输入的维度
        dim_hidden: int,  # MLP的隐藏层维度
        dim_key: int,  # 压缩记忆中键的维度
        dim_value: int,  # 压缩记忆中值的维度
        num_heads: int,  # 压缩记忆中的注意力头数
        activation: str,  # MLP使用的激活函数名称
        segment_len: int,  # 压缩记忆中的段长度
        update: str = "linear",  # 压缩记忆的记忆更新规则类型，默认为 "linear"
        causal: bool = False,  # 是否使用因果注意力掩码，默认为 False
        position_embedder: Optional[
            PositionEmbeddings
        ] = None,  # 位置嵌入模块，默认为 None
        init_state_learnable: bool = False,  # 压缩记忆的初始状态是否可学习，默认为 False
        dropout: float = 0.0,  # MLP的丢弃率，默认为 0.0
        **kwargs,
    ):
        """初始化模块。"""
        super(InfiniTransformer, self).__init__()  # 调用父类的初始化方法

        # 如果 sampling_factor 传递给 kwargs，则使用它，否则设置为 None
        sampling_factor = kwargs.get("sampling_factor", None)

        # 多头注意力
        self.attn = CompressiveMemory(
            dim_input=dim_input,
            dim_key=dim_key,
            dim_value=dim_value,
            num_heads=num_heads,
            segment_len=segment_len,
            sampling_factor=sampling_factor,
            update=update,
            causal=causal,
            position_embedder=position_embedder,
            init_state_learnable=init_state_learnable,
        )
        # MLP
        if activation not in ACTIVATIONS:
            raise ValueError(f"Invalid activation function: {activation}")
        if activation in ["swiglu", "geglu", "ffnglu", "ffngeglu", "ffnswiglu"]:
            act = ACTIVATIONS[activation](dim_hidden)
        else:
            act = ACTIVATIONS[activation]()
        self.mlp = nn.Sequential(
            nn.Linear(dim_input, dim_hidden),
            nn.Dropout(dropout),
            act,
            nn.Linear(dim_hidden, dim_input),
            nn.Dropout(dropout),
        )
        self.layer_norm = nn.LayerNorm(dim_input)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播。"""
        # 应用多头注意力，后跟 MLP 和层归一化与残差连接。
        x_ = self.attn(x)
        x_ = self.mlp(x_)

        return self.layer_norm(x_ + x)


class MoDInfiniTransformer(InfiniTransformer):
    """Mixture-of-Depths Infini-Transformer 层。"""

    def __init__(
        self,
        dim_input: int,  # 输入的维度
        dim_hidden: int,  # MLP的隐藏层维度
        dim_key: int,  # 压缩记忆中键的维度
        dim_value: int,  # 压缩记忆中值的维度
        num_heads: int,  # 压缩记忆中的注意力头数
        activation: str,  # MLP使用的激活函数名称
        segment_len: int,  # 压缩记忆中的段长度
        sampling_factor: int,  # Mixture-of-Depths 机制的采样因子
        update: str = "linear",  # 压缩记忆的记忆更新规则类型，默认为 "linear"
        causal: bool = False,  # 是否使用因果注意力掩码，默认为 False
        position_embedder: Optional[
            PositionEmbeddings
        ] = None,  # 位置嵌入模块，默认为 None
        init_state_learnable: bool = False,  # 压缩记忆的初始状态是否可学习，默认为 False
        dropout: float = 0.0,  # MLP的丢弃率，默认为 0.0
    ):
        """初始化模块。"""
        # 初始化普通的 InfiniTransformer，但将段长度减少了 sampling_factor
        super(MoDInfiniTransformer, self).__init__(
            dim_input=dim_input,
            dim_hidden=dim_hidden,
            dim_key=dim_key,
            dim_value=dim_value,
            num_heads=num_heads,
            activation=activation,
            segment_len=math.ceil(segment_len / sampling_factor),
            update=update,
            causal=causal,
            position_embedder=position_embedder,
            init_state_learnable=init_state_learnable,
            dropout=dropout,
            sampling_factor=sampling_factor,
        )

        # 记录额外的初始化参数用于前向传播
        self.segment_len = math.ceil(segment_len / sampling_factor)
        self.full_segment_len = segment_len
        self.sampling_factor = sampling_factor
        self.dim_input = dim_input

        # 采样时使用的投影层
        self.proj_sampling = nn.Linear(dim_input, 1)

    def forward(
        self, x: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        """前向传播包装器 -- 用于检查推理时是否处理每个观察值。"""
        if self.training:
            return self.forward_(x)
        else:
            # !!! 临时：每个样本可能具有不同的序列长度，导致 ragged 数组
            # !!!        当前的解决方案是单独处理每个样本并连接结果
            out = []
            # 循环遍历样本并产生每个样本的输出
            for ix in range(x.size(0)):
                sample_out, _, _ = self.forward_(x[ix : ix + 1, ...])
                out.append(sample_out)

            # 连接结果
            out = torch.cat(out, dim=0)

            return out, None, None

    def forward_(
        self, x: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        """前向传播。"""
        # 计算总段数和样本数
        batch_size, seq_len, _ = x.shape
        num_segments, rem = divmod(seq_len, self.full_segment_len)
        num_segments += 1 if rem > 0 else 0

        # 初始化采样掩码列表
        sample_masks = []

        # 使用线性嵌入进行采样得分
        sample_scores = self.proj_sampling(x).squeeze(-1)

        # 对于每个段，采样得分最高的令牌
        for seg_num in range(num_segments):
            # 计算段索引
            ix_lo = seg_num * self.full_segment_len
            ix_hi = ix_lo + self.full_segment_len

            if self.training:
                # 在训练期间，取得分最高的前 k 个令牌
                # 使用 sample_scores 对令牌索引进行排序
                sort_ixs = torch.argsort(
                    sample_scores[:, ix_lo:ix_hi], dim=1, descending=True
                )

                # 将令牌索引转换为二进制掩码
                sample_mask_seg = torch.zeros_like(
                    sample_scores[:, ix_lo:ix_hi], device=x.device
                )
                sample_mask_seg.scatter_(
                    dim=1, index=sort_ixs[:, : self.segment_len], value=1.0
                )
            else:
                # 在推理期间，取得分大于零的令牌
                sample_mask_seg = (sample_scores[:, ix_lo:ix_hi] > 0.0).float()

            sample_masks.append(sample_mask_seg)

        # 将段掩码连接成单个掩码
        sample_mask = torch.cat(sample_masks, dim=1).bool()

        # 根据采样掩码提取输入张量的子序列
        sample_shape = (batch_size, self.segment_len * num_segments, self.dim_input)
        x_ = x[sample_mask.unsqueeze(-1).repeat((1, 1, self.dim_input))].view(
            sample_shape
        )

        # 应用多头注意力到采样，然后应用 MLP
        x_ = self.attn(x_, sample_mask=sample_mask)
        x_ = self.mlp(x_)

        # 将注意力令牌的结果添加到结果中（相当于将非注意力令牌的结果设置为零）
        x[sample_mask.unsqueeze(-1).repeat((1, 1, self.dim_input))] += x_.view(-1)

        # 填充输出张量以匹配原始序列长度
        padding_mask = (
            torch.arange(x.size(1), device=x.device)[None, :]
            < sample_mask.view(batch_size, -1).sum(dim=1)[:, None]
        )
        x = x * padding_mask.unsqueeze(-1)

        # 展平采样得分和连接的前 k 个掩码用于辅助训练任务
        sample_scores = sample_scores.view((-1, 1))
        sample_mask = sample_mask.view((-1, 1)).float()

        return x, sample_mask, sample_scores


def demo_mod_infini_transformer():
    """
    演示 MoDInfiniTransformer 类的使用。
    """
    # 定义模型参数
    dim_input = 512
    dim_hidden = 2048
    dim_key = 64
    dim_value = 64
    num_heads = 8
    activation = "ffngeglu"
    segment_len = 2048
    sampling_factor = 8
    update = "linear"
    dropout = 0.1
    position_embedder = None

    # 定义批次维度
    seq_len = 4096
    batch_size = 2

    # 创建 MoDInfiniTransformer 层
    layer = MoDInfiniTransformer(
        dim_input=dim_input,
        dim_hidden=dim_hidden,
        dim_key=dim_key,
        dim_value=dim_value,
        num_heads=num_heads,
        activation=activation,
        segment_len=segment_len,
        sampling_factor=sampling_factor,
        update=update,
        dropout=dropout,
        position_embedder=position_embedder,
    )

    # 生成虚拟批次
    x = torch.randn(batch_size, seq_len, dim_input)

    # 测试训练时的输出
    layer.train()
    x_att, sample_mask, sample_scores_pred = layer(x)

    # 测试推理时的输出
    layer.eval()
    x_att, sample_mask, sample_scores_pred = layer(x)
